package main

import (
	"embed"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"

	"gopkg.in/yaml.v3"

	"gorm.io/driver/mysql"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

//go:embed templates/*
var templateFS embed.FS

type Config struct {
	DSN            string `yaml:"dsn"`
	DBType         string `yaml:"db_type"`
	ModelPackage   string `yaml:"model_package"`
	OutputDir      string `yaml:"output_dir"`
	ServicePackage string `yaml:"service_package"`
	APIPackage     string `yaml:"api_package"`
}

type Field struct {
	Name     string
	Type     string
	GormTag  string
	GormType string
	JsonTag  string
}

type TemplateData struct {
	PackageName   string
	StructName    string
	ModelName     string
	ModelImport   string
	ServiceImport string
	Fields        []Field
	TableName     string
	HasTime       bool
}

type column struct {
	Field string
	Type  string
	Key   string
	Extra string
}

func loadConfig(path string, cfg *Config) {
	if path == "" {
		return
	}

	data, err := os.ReadFile(path)
	if err != nil {
		fmt.Printf("Error reading config file: %v\n", err)
		os.Exit(1)
	}

	if err = yaml.Unmarshal(data, cfg); err != nil {
		fmt.Printf("Error parsing YAML config: %v\n", err)
		os.Exit(1)
	}
}

func makeDir(cfg *Config) {
	// 创建输出目录
	fmt.Println("Creating output directory...")
	outputDir := filepath.Clean(cfg.OutputDir)
	if err := os.MkdirAll(outputDir, 0700); err != nil {
		fmt.Printf("Failed to create output directory: %v\n", err)
		os.Exit(1)
	}

	modelDir := filepath.Join(cfg.OutputDir, "model")
	if err := os.MkdirAll(modelDir, 0700); err != nil {
		fmt.Printf("Failed to create model directory: %v\n", err)
		os.Exit(1)
	}

	repoDir := filepath.Join(cfg.OutputDir, "repository")
	if err := os.MkdirAll(repoDir, 0700); err != nil {
		fmt.Printf("Failed to create repository directory: %v\n", err)
		os.Exit(1)
	}

	serviceDir := filepath.Join(cfg.OutputDir, "service")
	if err := os.MkdirAll(serviceDir, 0700); err != nil {
		fmt.Printf("Failed to create service directory: %v\n", err)
		os.Exit(1)
	}

	apiDir := filepath.Join(cfg.OutputDir, "api")
	if err := os.MkdirAll(apiDir, 0700); err != nil {
		fmt.Printf("Failed to create api directory: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Directory '%s' created successfully\n", outputDir)
}

func connectDB(cfg *Config) *gorm.DB {
	var dialector gorm.Dialector
	switch cfg.DBType {
	case "mysql":
		dialector = mysql.Open(cfg.DSN)
	case "sqlite":
		dialector = sqlite.Open(cfg.DSN)
	default:
		fmt.Printf("Unsupported database type: %s\n", cfg.DBType)
		os.Exit(1)
	}
	db, err := gorm.Open(dialector)
	if err != nil {
		fmt.Printf("Failed to connect database: %v\n", err)
		os.Exit(1)
	}
	return db
}

func getTables(db *gorm.DB, dbType string) []string {
	var tables []string
	query := ""
	switch dbType {
	case "mysql":
		query = "SHOW TABLES"
	case "sqlite":
		query = "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
	case "postgres":
		query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"
	default:
		fmt.Printf("Unsupported database type: %s\n", dbType)
		os.Exit(1)
	}

	if err := db.Raw(query).Pluck("Tables", &tables).Error; err != nil {
		fmt.Printf("Failed to get tables: %v\n", err)
		os.Exit(1)
	}
	return tables
}

func generateModel(cfg *Config, table string) {
	db, err := gorm.Open(mysql.Open(cfg.DSN))
	if err != nil {
		fmt.Printf("Failed to connect to database: %v\n", err)
		os.Exit(1)
	}
	var columns []column
	db.Raw(fmt.Sprintf("DESCRIBE %s", table)).Scan(&columns)

	hasTime := false
	fields := make([]Field, 0)
	for _, col := range columns {
		fieldType, gormType := mapDBType(cfg.DBType, col.Type)
		if fieldType == "time.Time" {
			hasTime = true
		}
		gormParts := []string{fmt.Sprintf("column:%s", col.Field)}
		if col.Key == "PRI" {
			gormParts = append(gormParts, "primaryKey")
		}
		if strings.Contains(col.Extra, "auto_increment") {
			gormParts = append(gormParts, "autoIncrement:true")
		}
		fields = append(fields, Field{
			Name:     toCamelCase(col.Field),
			Type:     fieldType,
			GormTag:  strings.Join(gormParts, ";"),
			GormType: gormType,
			JsonTag:  toSnakeCase(col.Field),
		})
	}

	tmpl, err := template.ParseFS(templateFS, "templates/model.tmpl")
	if err != nil {
		fmt.Printf("Failed to parse template: %v\n", err)
		os.Exit(1)
	}
	tmpl = template.Must(tmpl, err)

	f, err := os.Create(fmt.Sprintf("%s/model/%s.go", cfg.OutputDir, table))
	if err != nil {
		fmt.Printf("Failed to create file: %v\n", err)
		os.Exit(1)
	}

	defer func(f *os.File) { _ = f.Close() }(f)

	err = tmpl.Execute(f, TemplateData{
		PackageName: cfg.ModelPackage,
		StructName:  toCamelCase(table),
		TableName:   table,
		Fields:      fields,
		HasTime:     hasTime,
	})
	if err != nil {
		fmt.Printf("Template execution failed: %v\n", err)
		os.Exit(1)
	}
}

func generateRepository(cfg *Config, table string) {
	tmpl := template.Must(template.ParseFS(templateFS, "templates/repository.tmpl"))

	f, err := os.Create(fmt.Sprintf("%s/repository/%s_repository.go", cfg.OutputDir, table))
	if err != nil {
		fmt.Printf("Failed to create file: %v\n", err)
		os.Exit(1)
	}
	defer func(f *os.File) {
		_ = f.Close()
	}(f)

	if err := tmpl.Execute(f, TemplateData{
		PackageName: cfg.ModelPackage,
		StructName:  toCamelCase(table),
		ModelName:   toCamelCase(table),
		ModelImport: cfg.ModelPackage,
	}); err != nil {
		fmt.Printf("Repository template execution failed: %v\n", err)
		os.Exit(1)
	}
}

func generateService(cfg *Config, table string) {
	tmpl := template.Must(template.ParseFS(templateFS, "templates/service.tmpl"))

	f, err := os.Create(fmt.Sprintf("%s/service/%s_service.go", cfg.OutputDir, table))
	if err != nil {
		fmt.Printf("Failed to create file: %v\n", err)
		os.Exit(1)
	}
	defer func(f *os.File) {
		_ = f.Close()
	}(f)

	if err = tmpl.Execute(f, TemplateData{
		PackageName: cfg.ServicePackage,
		StructName:  toCamelCase(table),
		ModelName:   toCamelCase(table),
		ModelImport: cfg.ModelPackage,
	}); err != nil {
		fmt.Printf("Service template execution failed: %v\n", err)
		os.Exit(1)
	}
}

func generateAPI(cfg *Config, table string) {
	tmpl := template.Must(template.ParseFS(templateFS, "templates/api.tmpl"))

	f, err := os.Create(fmt.Sprintf("%s/api/%s_api.go", cfg.OutputDir, table))
	if err != nil {
		fmt.Printf("Failed to create file: %v\n", err)
		os.Exit(1)
	}
	defer func(f *os.File) {
		_ = f.Close()
	}(f)

	if err = tmpl.Execute(f, TemplateData{
		PackageName:   cfg.APIPackage,
		StructName:    toCamelCase(table),
		ModelName:     toCamelCase(table),
		ModelImport:   fmt.Sprintf("\"%s\"", cfg.ModelPackage),
		ServiceImport: fmt.Sprintf("\"%s\"", cfg.ServicePackage),
	}); err != nil {
		fmt.Printf("API template execution failed: %v\n", err)
		os.Exit(1)
	}
}

func mapDBType(dbType string, columnType string) (string, string) {
	//fmt.Println("====================================<<<<< ", dbType, columnType)
	var gormType string
	switch dbType {
	case "sqlite":
		switch strings.ToLower(columnType) {
		case "integer":
			return "int", "INTEGER"
		case "text":
			return "string", "TEXT"
		case "real":
			return "float64", "REAL"
		case "blob":
			return "[]byte", "BLOB"
		default:
			return "string", "TEXT"
		}
	case "mysql":
		gormType = "type:" + strings.ToUpper(columnType)
		switch {
		case strings.HasPrefix(columnType, "enum"):
			return "string", gormType
		case strings.HasPrefix(columnType, "timestamp"), strings.HasPrefix(columnType, "datetime"):
			return "time.Time", gormType
		case strings.HasPrefix(columnType, "bigint"):
			return "int64", gormType
		case strings.HasPrefix(columnType, "integer"), strings.HasPrefix(columnType, "int"),
			strings.HasPrefix(columnType, "tinyint"), strings.HasPrefix(columnType, "smallint"):
			return "int", gormType
		case strings.HasPrefix(columnType, "varchar"), strings.HasPrefix(columnType, "text"),
			strings.HasPrefix(columnType, "char"), strings.HasPrefix(columnType, "time"):
			return "string", gormType
		case strings.HasPrefix(columnType, "tinyint(1)"):
			return "int32", gormType
		case strings.HasPrefix(columnType, "float"):
			return "float32", gormType
		case strings.HasPrefix(columnType, "decimal"):
			return "float64", gormType

		}
	case "postgres":
		switch {
		case strings.Contains(columnType, "text"):
			return "string", gormType
		case strings.Contains(columnType, "integer"):
			return "int", gormType
		case strings.Contains(columnType, "timestamp"):
			gormType = "type:TIMESTAMP"
			return "time.Time", gormType
		case strings.Contains(columnType, "boolean"):
			return "bool", gormType
		}
	}
	return "any", gormType
}

func toCamelCase(s string) string {
	parts := strings.Split(s, "_")
	for i := range parts {
		if parts[i] == "" {
			continue
		}
		parts[i] = strings.Title(strings.ToLower(parts[i]))
	}
	return strings.Join(parts, "")
}

func toSnakeCase(s string) string {
	var matchFirstCap = regexp.MustCompile("(.)([A-Z][a-z]+)")
	var matchAllCap = regexp.MustCompile("([a-z0-9])([A-Z])")
	snake := matchFirstCap.ReplaceAllString(s, "${1}_${2}")
	snake = matchAllCap.ReplaceAllString(snake, "${1}_${2}")
	return strings.ToLower(snake)
}

func main() {
	var cfg Config
	configPath := flag.String("c", "config.yaml", "Path to config file (YAML format)")
	flag.Parse()

	loadConfig(*configPath, &cfg)

	if cfg.DSN == "" || cfg.OutputDir == "" {
		fmt.Println("Error: dsn and output directory must be set in config file or command line parameters")
		os.Exit(1)
	}

	makeDir(&cfg)

	db := connectDB(&cfg)
	tables := getTables(db, cfg.DBType)

	for _, table := range tables {
		generateModel(&cfg, table)
		generateRepository(&cfg, table)
		generateService(&cfg, table)
		//generateAPI(&cfg, table)
	}
}
